#include "kpm_solver.h"
#include "kpm_dos_kernels.cu"
#include "kpm_ldos_kernels.cu"
#include "randomseed.hpp"
#include "cuda_runtime.h"
#include "stopwatch.h"
#include <fstream>
#include <time.h>
#include <mtrnd.h>

#define cutilSafeCall(err)     __cudaSafeCall	(err,__FILE__,__LINE__)

inline void __cudaSafeCall(cudaError err, const char *file, const int line )
{
    if( cudaSuccess != err) {
		printf( "%s(%i) : cudaSafeCall() Runtime API error : %s.\n",
                file, line, cudaGetErrorString( err));
        exit(-1);
    }
}

template<typename T>
T* solveDOSPlan_gpu(DOSPlan<T>* devicePlan)
{
	cudaSetDevice(devicePlan->device);

	 ELL<T>* H = devicePlan->H;
	 int h_size = H->rowsNum;
	 int elementsNum = H->rowsNum * H->colsNum;
	 int M = devicePlan->m;

	 printf("H_SIZE is: %d\n",H->rowsNum);
	 //cudaFuncSetCacheConfig(cukpmCreateRandR1<T, BLOCK_SIZE>, cudaFuncCachePreferL1);


	 //three vectors (d_R, d_R1, d_R2) are needed to perform the recursion.
	 T *d_R;
	 cutilSafeCall(cudaMalloc((void**)&d_R,sizeof(T) * h_size));

	 T *d_R1;
	 cutilSafeCall(cudaMalloc((void**)&d_R1,sizeof(T) * h_size));

	 T *d_R2;
	 cutilSafeCall(cudaMalloc((void**)&d_R2,sizeof(T) * h_size));

	 T *d_tempMu;
	 cutilSafeCall(cudaMalloc((void**)&d_tempMu,sizeof(T) * BLOCKS * M));
	 cutilSafeCall(cudaMemset(d_tempMu, 0, sizeof(T) * BLOCKS * M));

	 T *d_Mu;
	 cutilSafeCall(cudaMalloc((void**)&d_Mu,sizeof(T) * M));
	 cutilSafeCall(cudaMemset(d_Mu, 0, sizeof(T) * M));

	 //start to time
	 //Timer tm;
	 //startTimer(&tm);

	 ELL<T>* cudaH = (ELL<T>*)malloc(sizeof(ELL<T>));

	 cudaH->colsNum = H->colsNum;
	 cudaH->rowsNum = H->rowsNum;
	 //malloc memory for cudaH on device
	 cutilSafeCall(cudaMalloc((void**)&cudaH->elements,sizeof(T) * elementsNum));
	 cutilSafeCall(cudaMalloc((void**)&cudaH->cols,sizeof(int) * elementsNum));

	 //copy hamiltonian matrix to device
	 cutilSafeCall(cudaMemcpy(cudaH->elements,H->elements,sizeof(T) * elementsNum,cudaMemcpyHostToDevice));
	 cutilSafeCall(cudaMemcpy(cudaH->cols,H->cols,sizeof(int)* elementsNum ,cudaMemcpyHostToDevice));


	 //copy struct cudaH to device
	 ELL<T>* d_H;
	 cutilSafeCall(cudaMalloc((void**)&d_H,sizeof(ELL<T>)));
	 cutilSafeCall(cudaMemcpy(d_H,cudaH,sizeof(ELL<T>),cudaMemcpyHostToDevice));

	 //create random states for random generator
	 curandStateXORWOW_t* state;
	 cutilSafeCall(cudaMalloc((void**)&state,sizeof(curandStateXORWOW_t) * BLOCKS * BLOCK_SIZE));

	 //printf("cuda: %s\n",cudaGetErrorString(cudaGetLastError()));
	 int i =0,j=0;

	 int num_r = devicePlan->num_r;

	 cudaStream_t stream;
	 cudaStreamCreate(&stream);
	 printf("cuda: %s\n",cudaGetErrorString(cudaGetLastError()));
	 InitializeRand<<<BLOCKS,BLOCK_SIZE,1,stream>>>(state, RandomSeed::generate());
	 printf("cuda: %s\n",cudaGetErrorString(cudaGetLastError()));
	 //if(checkError(120)) exit(0);

	 int smemSize = BLOCK_SIZE * sizeof(T);


		StopWatchC timer = StopWatch::get(StopWatch::create());
	    timer.start();

	 for(i=0;i<num_r;i++){
		 cukpmCreateRandR1<T, BLOCK_SIZE><<<BLOCKS,BLOCK_SIZE,smemSize>>>
				 (state, d_R, d_R1, h_size, &d_tempMu[0]);
		 //if(checkError(126)) exit(0);

		 cukpmCreateR2<T, BLOCK_SIZE><<<BLOCKS, BLOCK_SIZE, smemSize>>>
				 (d_H, d_R1, d_R2, d_R, &d_tempMu[BLOCKS]);

		 for(j=2;j<M;j++){
			 cukpmCreateRn<T, BLOCK_SIZE><<<BLOCKS,BLOCK_SIZE, smemSize>>>
						 (d_H, d_R1, d_R2, d_R, &d_tempMu[BLOCKS * j]);
				 //if(checkError(134)) exit(0);
				 T *pTemp = d_R1;
				 d_R1=d_R2;
				 d_R2=pTemp;
			 }
	   }
	   printf("cuda: %s\n",cudaGetErrorString(cudaGetLastError()));
	   cukpmSumMu<T, BLOCKS><<<64, BLOCKS, smemSize>>>(d_tempMu, d_Mu, M, h_size * devicePlan->num_r);

	   cudaMemcpy(devicePlan->U,d_Mu,sizeof(T) * M,cudaMemcpyDeviceToHost);
	   timer.stop();
	   printf("Time used: %.3f\n", timer.getTime());

		time_t t;time(&t);
	   std::ofstream o_file;
	   o_file.open("time.txt", std::ios::app);
	   o_file<<std::endl<<"---------------------------"<<std::endl;
	   o_file<<"GPU H_SIZE: "<<devicePlan->H->rowsNum<<std::endl;
	   o_file<<timer.getTime()<<"ms";
	   o_file<<std::endl<<ctime(&t)<<std::endl;
	   o_file<<std::endl<<"---------------------------"<<std::endl;
	   o_file.close();

	   //printf("cuda: %s\n",cudaGetErrorString(cudaGetLastError()));



	 T* dd = (T*)malloc(sizeof(T) * BLOCKS * M);
	 cudaMemcpy(dd, d_tempMu, sizeof(T) * BLOCKS * M,cudaMemcpyDeviceToHost);
	 //endTimer(&tm);
	 //printf("TimeUsed: %.4f",getTimeDiff(&tm));




	printf("\nINFO:*********************\n");
	int ii;
	for(ii=0;ii<M;ii++)
	{
		printf("%.3f ", devicePlan->U[ii]);
	}
	printf("\n*********************\n");


	if(cudaH->elements == NULL)
		printf("malloc memory failded!");

	  cudaFree(cudaH->cols);
	  cudaFree(cudaH->elements);
	  cudaFree(d_H);
	  cudaFree(d_R);
	  cudaFree(d_R1);
	  cudaFree(d_R2);
	  cudaFree(d_Mu);
	  cudaFree(d_tempMu);
	  cudaFree(state);

	  free(cudaH);

	  return NULL;
}

template<typename T, int isMean>
T* solveLDOSPlan_gpu(LDOSPlan<T>* devicePlan){

	 cudaSetDevice(devicePlan->device);

	 ELL<T>* H = devicePlan->H;
	 int h_size = H->rowsNum;
	 int elementsNum = H->rowsNum * H->colsNum;
	 int M = devicePlan->m;

	 printf("H_SIZE is: %d\n",H->rowsNum);
	 //cudaFuncSetCacheConfig(cukpmCreateRandR1<T, BLOCK_SIZE>, cudaFuncCachePreferL1);


	 //three vectors (d_R0, d_R1) are needed to perform the recursion.

	 T *d_R0;
	 cutilSafeCall(cudaMalloc((void**)&d_R0,sizeof(T) * h_size));

	 T *d_R1;
	 cutilSafeCall(cudaMalloc((void**)&d_R1,sizeof(T) * h_size));

	 T *d_Mu;
	 cutilSafeCall(cudaMalloc((void**)&d_Mu,sizeof(T) *  devicePlan->m));
	 cutilSafeCall(cudaMemset(d_Mu, 0, sizeof(T) *  devicePlan->m));

	 T *d_LDOS;
	 if(isMean == 0){
		 cutilSafeCall(cudaMalloc((void**)&d_LDOS,sizeof(T) * devicePlan->num_s * devicePlan->num_e));
		 cutilSafeCall(cudaMemset(d_LDOS, 0, sizeof(T) * devicePlan->num_s * devicePlan->num_e));
	 }else{
		 cutilSafeCall(cudaMalloc((void**)&d_LDOS,sizeof(T) * 2 * devicePlan->num_e));
		 cutilSafeCall(cudaMemset(d_LDOS, 0, sizeof(T) * 2 * devicePlan->num_e));
	 }

	 T* d_E;
	 cutilSafeCall(cudaMalloc((void**)&d_E, sizeof(T)*devicePlan->num_e));
	 cutilSafeCall(cudaMemcpy(d_E, devicePlan->E, sizeof(T)*devicePlan->num_e, cudaMemcpyHostToDevice));


	 //start to time
	 //Timer tm;
	 //startTimer(&tm);

	 ELL<T>* cudaH = (ELL<T>*)malloc(sizeof(ELL<T>));

	 cudaH->colsNum = H->colsNum;
	 cudaH->rowsNum = H->rowsNum;
	 //malloc memory for cudaH on device
	 cutilSafeCall(cudaMalloc((void**)&cudaH->elements,sizeof(T) * elementsNum));
	 cutilSafeCall(cudaMalloc((void**)&cudaH->cols,sizeof(int) * elementsNum));

	 //copy hamiltonian matrix to device
	 cutilSafeCall(cudaMemcpy(cudaH->elements,H->elements,sizeof(T) * elementsNum,cudaMemcpyHostToDevice));
	 cutilSafeCall(cudaMemcpy(cudaH->cols,H->cols,sizeof(int)* elementsNum ,cudaMemcpyHostToDevice));


	 //copy struct cudaH to device
	 ELL<T>* d_H;
	 cutilSafeCall(cudaMalloc((void**)&d_H, sizeof(ELL<T>)));
	 cutilSafeCall(cudaMemcpy(d_H, cudaH, sizeof(ELL<T>), cudaMemcpyHostToDevice));


	 //printf("cuda: %s\n",cudaGetErrorString(cudaGetLastError()));
	 int i =0,j=0;

	 //number of samples of |i>
	 int num_s = devicePlan->num_s;

	 printf("cuda: %s\n",cudaGetErrorString(cudaGetLastError()));

	 //if(checkError(120)) exit(0);

	 MT::MersenneTwist mt;
	 mt.init_genrand(RandomSeed::generate());

	 StopWatchC timer = StopWatch::get(StopWatch::create());
	 timer.start();

	 //if(isMean != 0)
		// LDOS::cukpmSetOneToVector<T><<<BLOCKS,BLOCK_SIZE>>>(&d_LDOS[devicePlan->num_e], devicePlan->num_e);

	 for(i=0;i<num_s;i++){
		 unsigned int r = mt.genrand_int32() % h_size;

		 LDOS::cukpmCreateR0<T><<<BLOCKS,BLOCK_SIZE>>>(r, d_R0, h_size, d_Mu);

		 LDOS::cukpmCreateR1<T><<<BLOCKS,BLOCK_SIZE>>>(d_H, d_R0, d_R1, d_Mu);

		 for(j=2;j<M;j++){
			 LDOS::cukpmCreateRn<T><<<BLOCKS,BLOCK_SIZE>>>
					 (d_H, d_R0, d_R1, d_Mu, r, j);
				 //if(checkError(134)) exit(0);
				 T *pTemp = d_R0;
				 d_R0=d_R1;
				 d_R1=pTemp;
			 }
		 LDOS::cukpmCreateLastMu<T><<<BLOCKS, BLOCK_SIZE>>>(d_Mu, M, d_R1,r);

		 if(isMean == 0)
			 LDOS::cukpmCreateDOS<T><<<BLOCKS, BLOCK_SIZE>>>(d_Mu, M, d_E, &d_LDOS[i * devicePlan->num_e], devicePlan->num_e);
		 else
			 LDOS::cukpmCreateMeanDOS<T><<<BLOCKS, BLOCK_SIZE>>>(d_Mu, M, d_E, &d_LDOS[0], &d_LDOS[devicePlan->num_e], devicePlan->num_e);

	   }

	   timer.stop();
	   printf("Time used: %.5f\n", timer.getTime());

	   if(isMean == 0)
		   //copy ldos values back to cpu side.
		   cutilSafeCall(cudaMemcpy(devicePlan->LDOS, d_LDOS,
				   sizeof(T) * devicePlan->num_s * devicePlan->num_e,
				   cudaMemcpyDeviceToHost));
	   else {
		   //get arithmetic mean and geometry mean.
		   LDOS::cukpmGetMean<T><<<BLOCKS, BLOCK_SIZE>>>(&d_LDOS[0], &d_LDOS[devicePlan->num_e], devicePlan->num_e, devicePlan->num_s);

		   //copy mean values back to cpu side.
		   cutilSafeCall(cudaMemcpy(devicePlan->LDOS, d_LDOS,
				   sizeof(T) * 2 * devicePlan->num_e,
				   cudaMemcpyDeviceToHost));
	   }

	   time_t t;time(&t);
	   std::ofstream o_file;
	   o_file.open("time.txt", std::ios::app);
	   o_file<<std::endl<<"---------------------------"<<std::endl;
	   o_file<<"GPU H_SIZE: "<<devicePlan->H->rowsNum<<std::endl;
	   o_file<<timer.getTime()<<"ms";
	   o_file<<std::endl<<ctime(&t)<<std::endl;
	   o_file<<std::endl<<"---------------------------"<<std::endl;
	   o_file.close();

	   for(int i=2000;i<2100;i++){
		   printf("%.6f ", devicePlan->LDOS[i]);
	   }
/*
	  cudaFree(cudaH->cols);
	  cudaFree(cudaH->elements);
	  cudaFree(d_H);

	  cudaFree(d_R0);
	  cudaFree(d_R1);
	  cudaFree(d_Mu);
	  cudaFree(d_LDOS);
	  cudaFree(d_E);

	  free(cudaH);
*/
	  return NULL;
}




template float* solveDOSPlan_gpu<float>(DOSPlan<float>* devicePlan);
template double* solveDOSPlan_gpu<double>(DOSPlan<double>* devicePlan);

template float* solveLDOSPlan_gpu<float, 1>(LDOSPlan<float>* devicePlan);
template double* solveLDOSPlan_gpu<double, 1>(LDOSPlan<double>* devicePlan);

template float* solveLDOSPlan_gpu<float, 0>(LDOSPlan<float>* devicePlan);
template double* solveLDOSPlan_gpu<double, 0>(LDOSPlan<double>* devicePlan);
